;; collection/tree.scm - (c) rohan drape, 2000-2005

;; Evaluates to '#t' iff `tree' is a leaf node of a tree.

(define (tree-leaf? tree)
  (not (list? tree)))

;; Evaluates to the result of applying the commutative binary function
;; `procedure' throughout the leaves of `tree', using `nil' as the
;; identity element.  Obvious candidates for `procedure' are '+', '*',
;; 'min' and 'max'.

(define (tree-reduce procedure nil tree)
  (if (null? tree)
      nil
      (let ((left (car tree))
	    (right (cdr tree)))
	(procedure (if (list? left)
		       (tree-reduce procedure nil left)
		       (procedure nil left))
		   (tree-reduce procedure nil right)))))

;; Traverse `tree' and construct a list of the results of applying
;; `procedure' to each leaf element.

(define (tree-walk procedure tree)
  (if (null? tree)
      '()
      (let ((left (car tree))
	    (right (cdr tree)))
	(append (if (list? left)
		    (tree-walk procedure left)
		    (cons (procedure left) '()))
		(tree-walk procedure right)))))

;; Traverse `tree' and construct a list of all the leaf elements.

(define (tree-leaves tree)
  (tree-walk (lambda (x) x) tree))

;; Traverse `tree' and construct a new tree by applying `procedure' to
;; each leaf element.

(define (tree-map procedure tree)
  (if (null? tree)
      '()
      (let ((left (car tree))
	    (right (cdr tree)))
	(cons (if (list? left)
		  (tree-map procedure left)
		  (procedure left))
	      (tree-map procedure right)))))

;; Map variant where each application has the depth of the leaf as an
;; initial argument.

(define (tree-map/depth depth procedure tree)
  (if (null? tree)
      '()
      (let ((left (car tree))
	    (right (cdr tree)))
	(cons (if (list? left)
		  (tree-map/depth (+ depth 1) procedure left)
		  (procedure depth left))
	      (tree-map/depth depth procedure right)))))

;; Map variant that applies `f' at all subtrees and leaves of `tree'.

(define (tree-map* f tree)
  (if (null? tree)
      '()
      (let ((left (car tree))
	    (right (cdr tree)))
	(cons (if (list? left)
		  (f (tree-map* f left))
		  (f left))
	      (f (tree-map* f right))))))

;; Traverse `tree' and in left to right order applying `procedure' to
;; each leaf element.

(define (tree-for-each procedure tree)
  (if (null? tree)
      '()
      (let ((left (car tree))
	    (right (cdr tree)))
	(if (list? left)
	    (tree-for-each procedure left)
	    (procedure left))
	(tree-for-each procedure right))))

;; Evaluates to the first sub-tree of `tree' that `predicate'
;; evaluates to '#t' over, or '#f' if no such sub-tree exists.  Note
;; that predicate is applied both to sub-trees and to leaf elements.

(define (tree-find predicate tree)
  (call-with-current-continuation
   (lambda (return)
     (define (search t)
       (if (null? t)
	   #f
	   (let ((left (car t))
		 (right (cdr t)))
	     (cond ((predicate left) (return left))
		   ((predicate right) (return right))
		   (else (begin 
			   (if (list? left)
			       (search left))
			   (search right)))))))
     (search tree))))

;; Variant that only evaluates predicate at leaf elements.

(define (tree-find-leaf predicate tree)
  (tree-find (lambda (tree)
	       (and (tree-leaf? tree)
		    (predicate tree)))
	     tree))

;; Variant that only evaluates predicate at non-leaf elements.

(define (tree-find-subtree predicate tree)
  (tree-find (and? list? predicate) tree))

;; Variant that finds all sub-trees at `tree' that satisfy `predicate'
;; and collects them in a list.  Note that sub-trees that are
;; collected are not descended.

(define (tree-collect predicate tree)
  (define (search r t)
    (if (not (list? t))
	'()
	(if (null? t)
	    r
	    (let ((left (car t))
		  (right (cdr t)))
	      (cond ((predicate left) (search (cons left r) right))
		    ((predicate right) (search (cons right r) left))
		    (else (append r (search '() left) (search '() right))))))))
  (search '() tree))

;; Traverse the tree `t' discarding all elements that do not satisfy
;; the predicate `p'.  Note that `p' is applied to both sub-trees and
;; leaf nodes.

(define (tree-filter p t)
  (if (not (list? t))
      t
      (if (null? t)
	  '()
	  (let ((left (car t))
		(right (cdr t)))
	    (if (p left)
		(cons (tree-filter p left) (tree-filter p right))
		(tree-filter p right))))))

;; Variant that applies the predicate only to leaf nodes.

(define (tree-filter-leaves p t)
  (tree-filter (lambda (e) (or (list? e) (p e))) t))

;; Variant that deletes nodes that satisy the predicate.

(define (tree-remove p t) 
  (tree-filter (lambda (x) (not (p x))) t))

;; Variant that applies the predicate only to leaf nodes.

(define (tree-remove-leaves p t) 
  (tree-filter (lambda (e) (or (list? e) (not (p e)))) t))

;; The basic binary tree operator, where `p' is a procedure to combine
;; two leaves of the trees `a' and `b'.  When `a' and `b' are lists of
;; unequal length, excess elements on the longer list are discarded.
;; This is the binary operator form of 'tree-map'.

(define (tree-combine p a b)
  ;; There are four cases: a and b are both lists, a is a list, b is a
  ;; list, neither are lists.
  (cond ((and (list? a) (list? b))
	 (map (lambda (aa bb) (tree-combine p aa bb)) a b))
	((list? a)
	 (map (lambda (aa) (tree-combine p aa b)) a))
	((list? b)
	 (map (lambda (bb) (tree-combine p a bb)) b))
	(else
	 (p a b))))

;; Traverse the tree `s' and replace all instances of a list with a
;; vector having the same elements.  Evaluates to the new structure.

(define (tree->vector s)
  (cond ((list? s)
	 (tree->vector (list->vector s)))
	((vector? s)
	 (do ((i 0 (+ i 1)))
	     ((= i (vector-length s)) s)
	   (vector-set! s i (tree->vector (vector-ref s i)))))
	(else
	 s)))

;; The clauses are *not* the same as the subtrees.  (2 3) is a subtree
;; of (1 2 3) but not a clause.

(define (tree-clauses tree)
  (append-map (lambda (e) 
		(if (list? e) 
		    (append (tree-clauses e) (list e))
		    (list e)))
	      tree))

;; Variant that only evaluates predicate at clauses.

(define (tree-find-clause predicate tree)
  (find predicate (tree-clauses tree)))

;; Make a tree that has the same shape as `t' with the elements of
;; `l'.

(define (list->tree l t)
  (tree-map (lambda (e) 
	      (let ((n (car l))) 
		(set! l (cdr l)) 
		n))
	    t))

